package tree;

import java.util.TreeSet;

public class RedBlackTree {
    static class RedBlackTreeNode extends Node{

        boolean isRed;

        RedBlackTreeNode parent;

        public RedBlackTreeNode(int value) {
            super(value);
            isRed = true;
        }

        /**
         * 以前序顺序打印该节点及其子树的内容
         * @param node  节点
         */
        public static void prePrintTree(RedBlackTreeNode node) {
            if (node == null) {
                return;
            }
            System.out.print((node.isRed ? "R" : "B") + node.value + " ");
            prePrintTree((RedBlackTreeNode) node.left);
            prePrintTree((RedBlackTreeNode) node.right);
        }

        /**
         * 以中序顺序打印该节点及其子树的内容
         * @param node  节点
         */
        public static void midPrintTree(RedBlackTreeNode node) {
            if (node == null) {
                return;
            }
            midPrintTree((RedBlackTreeNode) node.left);
            System.out.print((node.isRed ? "R" : "B") + node.value + " ");
            midPrintTree((RedBlackTreeNode) node.right);
        }
    }

    RedBlackTreeNode root;

    /**
     * 获取当前节点的兄弟节点
     * @param node  当前节点
     * @return  返回其兄弟节点
     */
    private static RedBlackTreeNode getBrother(RedBlackTreeNode node) {
        if (node.parent.left == node) {
            return (RedBlackTreeNode) node.parent.right;
        } else {
            return (RedBlackTreeNode) node.parent.left;
        }
    }

    /**
     * 连接某两个节点
     * @param node1 父节点
     * @param node2 子节点
     */
    private static void connect(RedBlackTreeNode node1, RedBlackTreeNode node2) {
        node1.parent = node2;
        if (node1.value < node2.value) {
            node2.left = node1;
        } else {
            node2.right = node1;
        }
    }

    /**
     * 插入某个值
     * @param num   值
     */
    public void insert(int num) {
        // 如果根节点为空则直接置为该值并设为黑色
        if (this.root == null) {
            root = new RedBlackTreeNode(num);
            root.isRed = false;
            return;
        }

        // 找到插入点的父节点
        RedBlackTreeNode p = root;
        boolean flag = false;
        while (!flag) {
            switch (Integer.compare(p.value, num)){
                case 1: {
                    if (p.left != null) {
                        p = (RedBlackTreeNode) p.left;
                    } else {
                        flag = true;
                    }
                    break;
                }
                case -1: {
                    if (p.right != null) {
                        p = (RedBlackTreeNode) p.right;
                    } else {
                        flag = true;
                    }
                    break;
                }
                case 0: {
                    return;     // 找到重复元素直接返回
                }
            }
        }

        // 先插入节点
        RedBlackTreeNode newNode = new RedBlackTreeNode(num);
        connect(newNode, p);
        if (p.isRed) {
            // 父节点是红色节点
            // 再获取叔叔节点
            RedBlackTreeNode uncle = getBrother(p);

            // 判断叔叔节点是否为红色节点
            if (uncle != null && uncle.isRed) {
                balanceWhenUncleIsRed(uncle, newNode);
            } else {
                balanceWhenUncleIsBlack(newNode);
            }
        }
    }

    /**
     * 当插入节点的叔叔节点为红色节点时维持平衡
     * @param uncle 叔叔节点
     * @param node  插入节点
     */
    private void balanceWhenUncleIsRed(RedBlackTreeNode uncle, RedBlackTreeNode node) {
        uncle.isRed = false;
        node.parent.isRed = false;
        node.parent.parent.isRed = true;

        // 向上递归看会不会有重复的红色节点出现
        RedBlackTreeNode q = node.parent.parent;
        if (q.parent == null) {
            q.isRed = false;
        } else if (q.parent.isRed) {
            RedBlackTreeNode newUncle = getBrother(q.parent);
            if (newUncle != null && newUncle.isRed) {
                balanceWhenUncleIsRed(newUncle, q);
            } else {
                balanceWhenUncleIsBlack(q);
            }
        }

    }

    /**
     * 节点左旋
     * @param node  节点
     */
    private static void leftRound(RedBlackTreeNode node) {
        RedBlackTreeNode tmp = (RedBlackTreeNode) node.right;
        node.right = tmp.left;
        tmp.left = node;
        tmp.parent = node.parent;
        node.parent = tmp;
        if (node.right != null) {
            ((RedBlackTreeNode) node.right).parent = node;
        }
        if (tmp.parent != null) {
            if (tmp.parent.left == node) {
                tmp.parent.left = tmp;
            } else {
                tmp.parent.right = tmp;
            }
        }
    }

    /**
     * 节点右旋
     * @param node  节点
     */
    private static void rightRound(RedBlackTreeNode node) {
        RedBlackTreeNode tmp = (RedBlackTreeNode) node.left;
        node.left = tmp.right;
        tmp.right = node;
        tmp.parent = node.parent;
        node.parent = tmp;
        if (node.left != null) {
            ((RedBlackTreeNode) node.left).parent = node;
        }
        if (tmp.parent != null) {
            if (tmp.parent.left == node) {
                tmp.parent.left = tmp;
            } else {
                tmp.parent.right = tmp;
            }
        }
    }

    /**
     * 当插入节点的叔叔节点为黑节点时维持平衡
     * @param node  插入节点
     */
    private void balanceWhenUncleIsBlack(RedBlackTreeNode node) {
        node.parent.isRed = false;
        node.parent.parent.isRed = true;
        if (node.parent.left == node) {
            if (node.parent.parent.left == node.parent) {
                // 左左
                rightRound(node.parent.parent);
            } else {
                // 左右
                leftRound(node.parent);
                rightRound(node.parent.parent);
            }
        } else {
            if (node.parent.parent.left == node.parent) {
                // 右左
                rightRound(node.parent);
                leftRound(node.parent.parent);
            } else {
                // 右右
                leftRound(node.parent.parent);
            }
        }

        // 原祖父节点为根节点则更改根节点为父节点
        if (getBrother(node) == root) {
            root = node.parent;
        }
    }

    /**
     * 删除某个值对应的节点
     * @param num   值
     * @return  删除成功返回true，反之返回false
     */
    public boolean remove(int num) {
        RedBlackTreeNode p = root;
        while (p != null) {
            switch (Integer.compare(p.value, num)) {
                case 1: {
                    p = (RedBlackTreeNode) p.left;
                    break;
                }
                case -1: {
                    p = (RedBlackTreeNode) p.right;
                    break;
                }
                case 0: {
                    removeNode(p);
                    return true;
                }
                default: {
                    break;
                }
            }
        }
        return false;
    }

    /**
     * 删除节点
     * @param p 要删除的节点
     */
    private void removeNode(RedBlackTreeNode p) {
        if (p.left == null && p.right == null) {
            // 如果是根节点直接删除
            if (p == root) {
                root = null;
                return;
            }
            RedBlackTreeNode bro = getBrother(p);
            if (p.parent.left == p) {
                p.parent.left = null;
            } else {
                p.parent.right = null;
            }
            if (!p.isRed) {
                balanceWhenDelete(bro);
            }
        } else if (p.left != null && p.right != null) {
            // 被删除节点有两个儿子，则找到前置节点替换当前节点然后删除前置节点
            RedBlackTreeNode q = (RedBlackTreeNode) p.left;
            while (q.right != null) {
                q = (RedBlackTreeNode) q.right;
            }
            p.value = q.value;
            removeNode(q);
        } else {
            // 被删除节点只有一个儿子，则将自己的值设置为儿子的值然后删除儿子节点
            RedBlackTreeNode q = (RedBlackTreeNode) (p.left == null ? p.right : p.left);
            p.value = q.value;
            removeNode(q);
        }
    }

    /**
     * 用于在删除节点后维持平衡
     * @param bro   删除节点的兄弟节点
     */
    private void balanceWhenDelete(RedBlackTreeNode bro) {
        RedBlackTreeNode parent = bro.parent;
        // 1. 兄黑
        if (!bro.isRed) {
            RedBlackTreeNode nephew1, nephew2;  // 对侄和顺侄
            if (parent.left == bro) {
               nephew1 = (RedBlackTreeNode) bro.left;
               nephew2 = (RedBlackTreeNode) bro.right;
            } else {
                nephew1 = (RedBlackTreeNode) bro.right;
                nephew2 = (RedBlackTreeNode) bro.left;
            }

            // 1.1. 对侄红
            if (nephew1 != null && nephew1.isRed) {
                // 父兄交替旋转，父红兄黑
                if (parent.left == bro) {
                    rightRound(parent);
                } else {
                    leftRound(parent);
                }

                if (root == parent) {
                    root = bro;
                }

                bro.isRed = false;
                parent.isRed = true;
            } else if (nephew2 != null && nephew2.isRed) {
                // 1.2. 顺侄红
                // 兄侄交替旋转，调换颜色，随后按照对侄红处理
                if (bro.left == nephew2) {
                    rightRound(bro);
                } else {
                    leftRound(bro);
                }

                bro.isRed = true;
                nephew2.isRed = false;

                // 按照对侄红处理
                balanceWhenDelete(nephew2);
            } else {
                // 1.3. 双侄黑
                boolean flag = !parent.isRed;

                // 父黑兄红
                parent.isRed = false;
                bro.isRed = true;

                // 如果父亲节点本来就是黑的话就向上递归
                if (flag) {
                    if (parent.parent != null) {
                        RedBlackTreeNode newBro = getBrother(parent);
                        if (newBro != null ) {
                            balanceWhenDelete(newBro);
                        }
                    }
                }
            }
        } else{
            // 2. 兄红
            RedBlackTreeNode nephew;
            // 父兄交替旋转，父红兄黑
            if (parent.left == bro) {
                nephew = (RedBlackTreeNode) bro.right;
                rightRound(parent);
            } else {
                nephew = (RedBlackTreeNode) bro.left;
                leftRound(parent);
            }

            if (root == parent) {
                root = bro;
            }

            parent.isRed = true;
            bro.isRed = false;

            // 按照兄黑进行处理
            balanceWhenDelete(nephew);
        }
    }

    /**
     * 查询某个值是否在红黑树中
     * @param num   值
     * @return  存在返回true，反之返回false
     */
    public boolean contains(int num) {
        RedBlackTreeNode p = root;
        while (p != null) {
            switch (Integer.compare(p.value, num)) {
                case 1: {
                    p = (RedBlackTreeNode) p.left;
                    break;
                }
                case -1: {
                    p = (RedBlackTreeNode) p.right;
                    break;
                }
                case 0: {
                    return true;
                }
                default: {
                    break;
                }
            }
        }
        return false;
    }

    public static void main(String[] args) {
        RedBlackTree tree = new RedBlackTree();
        for (int i = 10; i <= 28; ++i) {
            tree.insert(i);
            System.out.print("前序遍历结果为： ");
            RedBlackTreeNode.prePrintTree(tree.root);
            System.out.println();
            System.out.print("中序遍历结果为： ");
            RedBlackTreeNode.midPrintTree(tree.root);
            System.out.println();
        }
        /*
            前序遍历结果为： B17 R13 B11 B10 B12 B15 B14 B16 R21 B19 B18 B20 B23 B22 R25 B24 B26 R27
            中序遍历结果为： B10 B11 B12 R13 B14 B15 B16 B17 B18 B19 B20 R21 B22 B23 B24 R25 B26 R27

            前序遍历结果为： B17 R13 B11 B10 B12 B15 B14 B16 R21 B19 B18 B20 B23 B22 R25 B24 B27 R26
            中序遍历结果为： B10 B11 B12 R13 B14 B15 B16 B17 B18 B19 B20 R21 B22 B23 B24 R25 R26 B27

            前序遍历结果为： B17 B15 B14 B16 R21 B19 B18 B20 B23 B22 R25 B24 B27 R26 R28
            中序遍历结果为： B14 B15 B16 B17 B18 B19 B20 R21 B22 B23 B24 R25 R26 B27 R28

            前序遍历结果为： B13 B11 B10 B12 R17 B15 B14 B16 B19 B18 R21 B20 B23 R22 R24
            中序遍历结果为： B10 B11 B12 B13 B14 B15 B16 R17 B18 B19 B20 R21 R22 B23 R24
         */

        for (int i = 28; i >= 10; --i) {
            boolean remove = tree.remove(i);
            System.out.print("前序遍历结果为： ");
            RedBlackTreeNode.prePrintTree(tree.root);
            System.out.println();
            System.out.print("中序遍历结果为： ");
            RedBlackTreeNode.midPrintTree(tree.root);
            System.out.println();
        }

    }
}
